import numpy as np
import matplotlib as mpl
from matplotlib import pyplot as plt

mpl.use("TkAgg")


def draw_stacked_bar(data: np.ndarray, labels_stack: list, labels_axis: list, filepath=None, sort=True, xlabel="xLabel",
                     title="Title", show=True):
    """
    data的结构：堆叠个数 * 柱形图个数

    """
    plt.figure(figsize=(15, 10))
    N = data.shape[1]
    sum_data = data.sum(axis=0)
    if sort:
        arg = sum_data.argsort()
        print(arg)
        data = data[:, arg]
        labels_axis = [labels_axis[i] for i in arg]
        print(labels_axis)
        sum_data = sum_data[arg]

    cum_data = np.concatenate((np.zeros_like(data[0:1]), np.cumsum(data, axis=0)[:-1]), axis=0)
    # menMeans = np.array([20, 35, 30, 35, 27])
    # womenMeans = [25, 32, 34, 20, 25]
    # nonMeans = [12, 123, 43, 23, 12]
    # # menStd = (2, 3, 4, 1, 2)
    # # womenStd = (3, 5, 2, 3, 3)
    ind = np.arange(N)  # the x locations for the groups
    height = 0.35  # the width of the bars: can also be len(x) sequence
    ps = [plt.barh(ind, data[i], height, left=cum_data[i]) for i in range(data.shape[0])]
    # p1 = plt.barh(ind, menMeans, width)
    # p2 = plt.barh(ind, womenMeans, width, left=menMeans)
    # p3 = plt.barh(ind, nonMeans, width, left=menMeans + womenMeans)
    plt.xlabel(xlabel)
    # plt.title(title)
    plt.yticks(ind, labels_axis)
    max_len = np.max(sum_data)
    plt.xlim(xmax=max_len * 1.2)
    # plt.legend([pi[0] for pi in ps], labels_stack, loc=4,
    # bbox_to_anchor=(1.01, 1.01)
    # )

    for i, rect in enumerate(ps[0]):
        print(sum_data)
        width = sum_data[i]
        plt.text(1.01 * width, rect.get_y() + rect.get_height() * 0.5, f"{width:.3f}", ha='left', va='center')
    plt.tight_layout()
    if filepath is not None:
        plt.savefig(filepath, format='svg')
    if show:
        plt.show()


if __name__ == '__main__':
    """
    data = np.array([[0.11820904, 0.02942409, 0.05007597, 0.06424092, 0.11119691, 0.05852617, 0.07122805, 0.05402752,
                      0.03735846, 0.12995429, 0.01875739, 0.03199697, 0.04192901, 0.01314893, 0.02445963, 0.09667244,
                      0.04879421],
                     [0.003115264797507791, 0.07609680489793375, 0.07353010653473491, 0.06231206785986022,
                       0.0533681202470474, 0.07411303631518686, 0.07565135195033974, 0.049712354749611457,
                       0.08749870495065573, 0.04447152362564982, 0.06108520989610738, 0.06906900522971658,
                       0.04254080255302469, 0.04411246166457868, 0.04199235094763321, 0.06269249047260027,
                       0.07863834330781196],
                     [0.00882751, 0.01797439, 0.01659151, 0.01335179, 0.05107547, 0.03180043, 0.01197872, 0.04276139,
                      0.04087487, 0.00688577, 0.01678849, 0.00450645, 0.02470153, 0.31747055, 0.31747055, 0.0461579,
                      0.03078269]])
    x = [0.19537494, 0.73541944, 0.06920562]
    data[0][0] = 0
    data[0] /= data[0].sum()
    data[2][0] = 0
    data[2] /= data[2].sum()
    print(data[0].sum())
    for i in range(17):
        for j in range(3):
            data[j][i] = data[j][i] * x[j]
    labels_axis = ['SDG' + str(i) for i in range(1, 18)]
    labels_stack = ['Target Importance Index', 'Target Structural Importance Index', 'Target Difficulty Index']
    draw_stacked_bar(data, labels_stack, labels_axis, xlabel='priority')

    """
    data = np.array([[0.0033112582781456984, 0.08121892758753008, 0.07747989460993818, 0.0662338830061796, 0.0543547436573119, 0.07801703676597553, 0.07912592463529344, 0.052510152897737715, 0.09244528052316187, 0.04843008191417537, 0.0033112582781456984, 0.07414779660330015, 0.0462696743522795, 0.04768973491376349, 0.04583132178574607, 0.06608244284421969, 0.08354058734709661]])
    labels_axis = ['SDG' + str(i) for i in range(1, 18)]
    labels_stack = ['Target Structural Importance Index']
    draw_stacked_bar(data, labels_stack, labels_axis, xlabel='Target Structural Importance Index', filepath='./Target Structural Importance Index.svg')
